'''
Use CTT to generate initial data on cartesian grid.
The range of the spatial coordinates is [0, L], with
symmetries at x=0 and robin boundary condition at x=L:
    $${\partial(ru)\over r}=0$$
'''
from itertools import product
import numpy as np


L = 10    
dx = 0.1  
J = 1
M = 1

X = np.mgrid[0:L:dx,0:L:dx,0:L:dx]
#r = np.sqrt(x*x+y*y+z*z)

# a term resulted from puncture method
cal_long_stuff = lambda u: 15.75*J*J*(x*x+y*y)/((u+1)*r+M)**8

u = np.zeros(X.shape[1:])

def update_du(du):
    v = np.zeros_like(du)
    prev = lambda idx: (idx-1 if idx>0 else 1)
    def dd(axis):
        A = tuple(val if i!=axis else v+1     for i,v in enumerate(idx))
        B = tuple(val if i!=axis else prev(v) for i,v in enumerate(idx))
        return ( u[A] + u[B] - 2*u[idx] ) / ( dx*dx )
    def d(axis):
        A = tuple(val if i!=axis else v+1     for i,v in enumerate(idx))
        B = tuple(val if i!=axis else prev(v) for i,v in enumerate(idx))
        return ( u[A] - u[B] ) / ( 2*dx )
    def d2(a,b):
        A = tuple(val+1 if i==a or i==b else val for i,v in enumerate(idx))
        B = tuple(val-1 if i==a or i==b else val for i,v in enumerate(idx))
        C = tuple(val+1 if i==a else (val-1 if i==b else val)
                                            for i,v in enumerate(idx))
        D = tuple(val-1 if i==a else (val+1 if i==b else val)
                                            for i,v in enumerate(idx))
        return (u[A] + u[B] - u[C] - u[D]) / (4*dx*dx)
    for idx in product(range(N-1), repeat=3):
        v[idx] = sum(dd(i) for i in (0,1,2))
    for i,j in product(range(N-1), repeat=2):
        idx = (N-1, i, j)
        for axis in (0,1,2):
            b = (axis+1)%3
            c = (axis+2)%3
            x = X[(slice(-1),)+idx]
            v[idx]=(u[idx] + (x[axis]**2+x[b]**2)*dd(b) +
                    (x[axis]**2+x[c]**2)*dd(c) +x[b]*d(b)*3 + x[c]*d(c)*3+
                    x[b]*x[c]*d2(b,c)*2) / x[axis]**2
            idx = idx[1:] + (idx[0],)
    for i in range(N-1):
        idx = (i,N-1,N-1)
        for axis in (0,1,2):
            #TODO: handle the grids at edges
            idx = idx[1:] + (idx[0],)
    #TODO: handle the grids at corners
               

